// Copyright (c) 2022 Presto Labs Pte. Ltd.
// Author: chunhui

#pragma once

#include <cmath>
#include <memory>
#include <queue>

#include <glog/logging.h>

#include "coin2/strategy/util/time_indexed_heap.h"

namespace coin2::strategy::util {

class TopPercentile {
 public:
  explicit TopPercentile(double percentile)
      : percentile_(percentile),
        min_heap_(TimeIndexedHeap::Less),
        max_heap_(TimeIndexedHeap::Greater) {
    CHECK(percentile > 0 && percentile < 1);
  }

  void Add(int index, double value) {
    if (index_map_.find(index) == index_map_.end()) {
      index_map_[index] = size_;
      size_++;
    }
    auto id = static_cast<int64_t>(index);
    if (max_heap_.Contains(id)) {
      max_heap_.Update(id, value);
      if (max_heap_.Top().second > min_heap_.Top().second) {
        MoveMaxTopToMinHeap();
      }
    } else if (min_heap_.Contains(id)) {
      min_heap_.Update(id, value);
      if (max_heap_.Top().second > min_heap_.Top().second) {
        MoveMinTopToMaxHeap();
      }
    } else if (!max_heap_.Empty() && max_heap_.Top().second > value) {
      max_heap_.Update(id, value);
    } else {
      min_heap_.Update(id, value);
    }

    // Balancing
    size_t total_size = Size();
    auto target_max_heap_size = static_cast<size_t>(std::floor(total_size * percentile_));

    while (max_heap_.Size() > target_max_heap_size) {
      MoveMaxTopToMinHeap();
    }

    while (max_heap_.Size() < target_max_heap_size) {
      MoveMinTopToMaxHeap();
    }
  }

  size_t Size() const { return max_heap_.Size() + min_heap_.Size(); }

  bool Get(double* value) {
    if (!min_heap_.Empty() && !max_heap_.Empty()) {
      size_t total_size = Size() - 1;
      double x1 = static_cast<double>(max_heap_.Size() - 1) / total_size;
      double x2 = static_cast<double>(max_heap_.Size()) / total_size;
      double y1 = max_heap_.Top().second;
      double y2 = min_heap_.Top().second;
      *value = y1 + (percentile_ - x1) * (y2 - y1) / (x2 - x1);
      return true;
    } else if (!min_heap_.Empty()) {
      *value = min_heap_.Top().second;
      return true;
    } else if (!max_heap_.Empty()) {
      *value = max_heap_.Top().second;
      return true;
    }
    return false;
  }

 private:
  void MoveMaxTopToMinHeap() {
    const auto& [idx, val] = max_heap_.Top();
    min_heap_.Add(idx, val);
    max_heap_.Pop();
  }

  void MoveMinTopToMaxHeap() {
    const auto& [idx, val] = min_heap_.Top();
    max_heap_.Add(idx, val);
    min_heap_.Pop();
  }

  double percentile_;
  std::unordered_map<int, int> index_map_;
  int size_{0};
  TimeIndexedHeap min_heap_;
  TimeIndexedHeap max_heap_;
};

}  // namespace coin2::strategy::util
